package tracks.controllers.methods.MCTS;


import java.util.*;

import core.game.Observation;
import core.game.StateObservation;
import ontology.Types;
import tools.ElapsedCpuTimer;
import tools.Utils;
import tools.Vector2d;
import tracks.controllers.methods.MCTS.Agent;
import tracks.controllers.methods.MCTS.Heuristic;

public class SingleTreeNode
{
    private static final double HUGE_NEGATIVE = -10000000;
    private static final double HUGE_POSITIVE = 10000000;
    public static double epsilon = 1e-6;
    public static double egreedyEpsilon = 0.05;
    public StateObservation state;
    public SingleTreeNode parent;
    public SingleTreeNode[] children;

    public double totValue;
    public int nVisits;
    public static Random m_rnd;
    private int m_depth;
    protected static double[] bounds = new double[]{Double.MAX_VALUE, -Double.MAX_VALUE};
    public SingleTreeNode(Random rnd) {this(null, null, rnd);}
    ArrayList<Types.ACTIONS> suggestActions = new ArrayList<>();

    public SingleTreeNode(StateObservation state, SingleTreeNode parent, Random rnd) {
        this.state = state;
        this.parent = parent;
        this.m_rnd = rnd;
        this.totValue = 0;
        if (parent!=null) {
            m_depth = parent.m_depth+1;
        } else {
            m_depth = 0;
        }

        if(this.state != null) {
            Vector2d avatarPos = state.getAvatarPosition();
            Vector2d avatarOri = state.getAvatarOrientation();
            ArrayList<Types.ACTIONS> totalActions = state.getAvailableActions();
            for (Types.ACTIONS a : totalActions) {
                if (a != Types.ACTIONS.ACTION_USE) {
                    StateObservation stCopy = state.copy();
                    stCopy.advance(a);
                    Vector2d newAvatarPos = stCopy.getAvatarPosition();
                    Vector2d newAvatarOri = stCopy.getAvatarOrientation();
                    if (newAvatarPos.equals(avatarPos) && avatarOri.equals(newAvatarOri)) continue;
                    this.suggestActions.add(a);
                }
                else {
                    int indexx = (int) (avatarPos.x / state.getBlockSize() + avatarOri.x);
                    int indexy = (int) (avatarPos.y / state.getBlockSize() + avatarOri.y);
                    int xbound = state.getObservationGrid().length;
                    int ybound = state.getObservationGrid()[0].length;
                    if (indexx < 0 || indexx >= xbound || indexy < 0 || indexy >= ybound) continue;
//                    if (state.getObservationGrid()[indexx][indexy].size() == 0) continue;
//                    this.suggestActions.add(a);
                    for (Observation obs: state.getObservationGrid()[indexx][indexy]) {
                        ArrayList<Integer> category4 = new ArrayList<>(Arrays.asList(0, 5, 6, 8, 9, 19));
                        ArrayList<Integer> category3 = new ArrayList<>(Arrays.asList(15,16,17));
                        if (obs.category == 4 && obs.itype!=2) {
                            this.suggestActions.add(a);
                            break;
                        }
                        if (obs.category == 3 && (category3.contains(obs.itype))) {
                            this.suggestActions.add(a);
                            break;
                        }
                    }
                }
            }
            if (this.suggestActions.size() == 0)
                this.suggestActions = new ArrayList<Types.ACTIONS>(Arrays.asList(Agent.actions));
        }

        this.children = new SingleTreeNode[this.suggestActions.size()];
    }

    public void mctsSearch(ElapsedCpuTimer elapsedTimer) {

        double avgTimeTaken = 0;
        double acumTimeTaken = 0;
        long remaining = elapsedTimer.remainingTimeMillis();
        int numIters = 0;

        int remainingLimit = 5;
        while(remaining > 2*avgTimeTaken && remaining > remainingLimit){
            ElapsedCpuTimer elapsedTimerIteration = new ElapsedCpuTimer();
            SingleTreeNode selected = treePolicy();
            double delta = selected.rollOut();
            backUp(selected, delta);

            numIters++;
            acumTimeTaken += (elapsedTimerIteration.elapsedMillis()) ;

            avgTimeTaken  = acumTimeTaken/numIters;
            remaining = elapsedTimer.remainingTimeMillis();
            //System.out.println(elapsedTimerIteration.elapsedMillis() + " --> " + acumTimeTaken + " (" + remaining + ")");
        }
        System.out.println("-- " + numIters + " -- ( " + avgTimeTaken + ")");
    }
    public SingleTreeNode treePolicy() {

        SingleTreeNode cur = this;

        while (!cur.state.isGameOver() && cur.m_depth < Agent.ROLLOUT_DEPTH)
        {
            if (cur.notFullyExpanded()) {
                return cur.expand();

            } else {
                SingleTreeNode next = cur.uct();
                //SingleTreeNode next = cur.egreedy();
                cur = next;
            }
        }

        return cur;
    }

    public SingleTreeNode expand() {

        int bestAction = 0;
        double bestValue = -1;

        for (int i = 0; i < children.length; i++) {
            double x = m_rnd.nextDouble();
            if (x > bestValue && children[i] == null) {
                bestAction = i;
                bestValue = x;
            }
        }

        StateObservation nextState = state.copy();
        nextState.advance(this.suggestActions.get(bestAction));

        SingleTreeNode tn = new SingleTreeNode(nextState, this, this.m_rnd);
        children[bestAction] = tn;
        return tn;
    }

    public SingleTreeNode uct() {

        SingleTreeNode selected = null;
        double bestValue = -Double.MAX_VALUE;
        for (SingleTreeNode child : this.children)
        {
            double hvVal = child.totValue;
            double childValue =  hvVal / (child.nVisits + this.epsilon);


//            childValue = Utils.normalise(childValue, bounds[0], bounds[1]);

            double uctValue = childValue +
                    Agent.K * Math.sqrt(Math.log(this.nVisits + 1) / (child.nVisits + this.epsilon));

            // small sampleRandom numbers: break ties in unexpanded nodes
            uctValue = Utils.noise(uctValue, this.epsilon, this.m_rnd.nextDouble());     //break ties randomly

            // small sampleRandom numbers: break ties in unexpanded nodes
            if (uctValue > bestValue) {
                selected = child;
                bestValue = uctValue;
            }
        }

        if (selected == null)
        {
            throw new RuntimeException("Warning! returning null: " + bestValue + " : " + this.children.length);
        }

        return selected;
    }

    public double rollOut()
    {
        StateObservation rollerState = state.copy();
        int thisDepth = this.m_depth;

//        System.out.println("Avatar at "+rollerState.getAvatarPosition().toString());
        while (!finishRollout(rollerState,thisDepth)) {

            int action = m_rnd.nextInt(this.suggestActions.size());
            if ((m_rnd.nextDouble() < 0.15) && suggestActions.contains(state.getAvatarLastAction()) && suggestActions.size()<4) rollerState.advance(state.getAvatarLastAction());
            else rollerState.advance(this.suggestActions.get(action));
//
            thisDepth++;
//            System.out.print(action+" ");
        }
//        System.out.println("____");

        Heuristic s = new Heuristic();
        double delta = s.evaluate(rollerState);

        if(delta < bounds[0])
            bounds[0] = delta;

        if(delta > bounds[1])
            bounds[1] = delta;

        return delta;
    }

    public boolean finishRollout(StateObservation rollerState, int depth)
    {
        if(depth >= Agent.ROLLOUT_DEPTH)      //rollout end condition.
            return true;

        if(rollerState.isGameOver())               //end of game
            return true;

        return false;
    }

    public void backUp(SingleTreeNode node, double result)
    {
        double target = result;
        SingleTreeNode n = node;
        while(n != null)
        {
            n.nVisits++;
            n.totValue += target;
            n = n.parent;
//            target = target / (1+Math.exp(-target/10));
//            target = target * 0.9;
            if (target < -1000000) target = 0;
            target = target*0.9;
        }
    }

    public Types.ACTIONS mostVisitedAction() {
        int selected = -1;
        double bestValue = -Double.MAX_VALUE;
        boolean allEqual = true;
        double first = -1;

        for (int i=0; i<children.length; i++) {

            if(children[i] != null)
            {
                if(first == -1)
                    first = children[i].nVisits;
                else if(first != children[i].nVisits)
                {
                    allEqual = false;
                }

                double childValue = children[i].nVisits;
                childValue = Utils.noise(childValue, this.epsilon, this.m_rnd.nextDouble());     //break ties randomly
                if (childValue > bestValue) {
                    bestValue = childValue;
                    selected = i;
                }
            }
        }

        if (selected == -1)
        {
            System.out.println("Unexpected selection!");
            selected = 0;
        }else if(allEqual)
        {
            //If all are equal, we opt to choose for the one with the best Q.
            selected = bestAction();
        }
        return this.suggestActions.get(selected);
    }

    public int bestAction()
    {
        int selected = -1;
        double bestValue = -Double.MAX_VALUE;

        for (int i=0; i<children.length; i++) {

            if(children[i] != null) {
                double childValue = children[i].totValue / (children[i].nVisits + this.epsilon);
                childValue = Utils.noise(childValue, this.epsilon, this.m_rnd.nextDouble());     //break ties randomly
                if (childValue > bestValue) {
                    bestValue = childValue;
                    selected = i;
                }
            }
        }

        if (selected == -1)
        {
            System.out.println("Unexpected selection!");
            selected = 0;
        }

        return selected;
    }


    public boolean notFullyExpanded() {
        for (SingleTreeNode tn : children) {
            if (tn == null) {
                return true;
            }
        }

        return false;
    }
}

